#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Sanitize YOLO label txt files to have exactly 5 columns per line:
  class x y w h

Drops any extra columns (e.g., confidence) and skips malformed lines.
Processes all .txt files under given directories recursively.

Usage:
  python scripts/fix_label_format.py datasets/trial/labels/train datasets/trial/labels/val
"""

from __future__ import annotations

import sys
from pathlib import Path


def sanitize_file(path: Path) -> bool:
    changed = False
    if not path.is_file() or path.suffix.lower() != ".txt":
        return False
    lines_in = path.read_text(encoding="utf-8", errors="ignore").splitlines()
    lines_out = []
    for ln in lines_in:
        parts = ln.strip().split()
        if len(parts) < 5:
            # skip malformed line
            changed = True
            continue
        kept = parts[:5]
        if len(parts) != 5:
            changed = True
        # basic numeric validation
        try:
            _ = int(float(kept[0]))
            float(kept[1]); float(kept[2]); float(kept[3]); float(kept[4])
        except Exception:
            changed = True
            continue
        lines_out.append(" ".join(kept))

    if lines_out:
        new_txt = "\n".join(lines_out) + "\n"
    else:
        new_txt = ""
    if changed:
        path.write_text(new_txt, encoding="utf-8")
    return changed


def sanitize_dir(root: Path) -> int:
    n = 0
    for p in root.rglob("*.txt"):
        if sanitize_file(p):
            n += 1
    return n


def main(argv: list[str]) -> None:
    if len(argv) < 2:
        print("Usage: python scripts/fix_label_format.py <dir> [<dir> ...]")
        sys.exit(2)
    total = 0
    for d in argv[1:]:
        root = Path(d)
        if not root.exists():
            print(f"Skip missing: {root}")
            continue
        cnt = sanitize_dir(root)
        total += cnt
        print(f"Sanitized {cnt} files under {root}")
    print(f"Total sanitized files: {total}")


if __name__ == "__main__":
    main(sys.argv)

